from ultralytics import YOLO
import os
import re
import torch
# 用detect权重 重新标注 pose的xywhn


# Load a model
model = YOLO("yolov8x.pt")  # pretrained YOLOv8n model
model = YOLO(r"runs\detect\detect_x_1\train6\weights\best.pt")  # pretrained YOLOv8n model
dataPath = r'data\pose4\train'
# dataPath = r'data\pose4\val'
image_path = os.path.join(dataPath, r'images')
label_path_old = os.path.join(dataPath, r'labels2')
label_path_new = os.path.join(r'data\pose_detect', r'train')
# label_path_new = os.path.join(r'data\pose_detect', r'val')

# Run batched inference on a list of images
results = model(source=r"data\pose4\train\images")  # return a list of Results objects
# results = model(source=r"data\pose4\val\images")  # return a list of Results objects
# results = model([r'data\pose4\train\images\000002.jpg'])  # return a list of Results objects
pattern = r"(?<=\S\s)\S+(?:\s\S+){3}(?=\s)"
# print(len(results))
# Process results list
for result in results:
    img_name = result.path.split('\\')[-1]
    # print(img_name)
    id = img_name[:-4]
    # print(id)
    boxes = result.boxes  # Boxes object for bounding box outputs
    xywhn = [torch.tensor(0.5),torch.tensor(0.5),torch.tensor(1),torch.tensor(1)]
    if(len(boxes.xywhn)!=0):
        xywhn = boxes.xywhn[0]
        
    xn = xywhn[0]
    yn = xywhn[1]
    wn = xywhn[2]
    hn = xywhn[3]

    label_path_old_f = open(os.path.join(label_path_old, id+'.txt'),'r')
    label_path_new_f = open(os.path.join(label_path_new, id+'.txt'),'w')
    replacement = f'{round(xn.item(), 6)} {round(yn.item(), 6)} {round(wn.item(), 6)} {round(hn.item(), 6)}'
    for line in label_path_old_f.readlines():
        new_text = re.sub(pattern, replacement, line, count=1)
        # print(line)
        # print(new_text)
        label_path_new_f.write(new_text)
    label_path_old_f.close()
    label_path_new_f.close()